from core.cuda import cuda, cuda_module
from core.function import Function
from implement.functions.conv.conv2d_gradw import Conv2DGradW
from utils.common import pair

from utils.functions_conv import im2col_array


class Conv2d(Function):
    def __init__(self, stride=1, pad=0):
        """
        Conv2d 操作的函数类。

        Args:
            stride (int or tuple): 卷积操作的步幅，可以是单个整数或包含两个整数的元组。
            pad (int or tuple): 卷积操作的填充，可以是单个整数或包含两个整数的元组。
        """
        super().__init__()
        self.stride = pair(stride)
        self.pad = pair(pad)

    def forward(self, x, W, b):
        """
        前向传播，执行 Conv2d 操作。

        Args:
            x (Variable): 输入变量。
            W (Variable): 卷积核变量。
            b (Variable or None): 偏置变量，可以为 None。

        Returns:
            Variable: Conv2d 操作的输出变量。
        """
        xp = cuda_module

        KH, KW = W.shape[2:]
        col = im2col_array(x, (KH, KW), self.stride, self.pad, to_matrix=False)

        col = cuda.to_array(col)
        W = cuda.to_array(W)
        y = xp.tensordot(col, W, ((1, 2, 3), (1, 2, 3)))
        if b is not None:
            b = cuda.to_array(b)
            y += b
        y = xp.rollaxis(y, 3, 1)
        return y

    def backward(self, gy):
        """
        反向传播，计算输入变量的梯度。

        Args:
            gy (Variable): 输出变量的梯度。

        Returns:
            tuple: 包含输入变量的梯度和卷积核梯度的元组。
        """
        from utils.functions_collect import deconv2d

        x, W, b = self.inputs
        # ==== gx ====
        gx = deconv2d(gy, W, b=None, stride=self.stride, pad=self.pad,
                      outsize=(x.shape[2], x.shape[3]))
        # ==== gW ====
        gW = Conv2DGradW(self)(x, gy)
        # ==== gb ====
        gb = None
        if b.data is not None:
            gb = gy.sum(axis=(0, 2, 3))
        return gx, gW, gb

    def conv2d(self, x, W, b=None):
        return self(x, W, b)
